from typing import *
import numpy as np

class HMM():

    def __init__(self, N: int, M: int):
        self.N = N
        self.M = M
        self.PI = np.ones((N, ), dtype=np.float)  #加1平滑
        self.A = np.ones((N, N), dtype=np.float)  #加1平滑
        self.B = np.ones((N, M), dtype=np.float)  #加1平滑

    def getPI(self) -> np.ndarray:
        """ 获取初始状态概率 R ^ N"""
        raise NotImplementedError()

    def getA(self) -> np.ndarray:
        """ 获取状态转移概率 R ^ (N * N)"""
        raise NotImplementedError()

    def getB(self) -> np.ndarray:
        """ 获取发射概率 R ^ (N * M) """
        raise NotImplementedError()


class TextSequence(HMM):
    def __init__(self, texts: List[List[str]], labels: List[List[int]], nClasses:int):
        self.texts = texts
        self.labels = labels

        # 构造词典
        dictSet = {"[OOV]": 0}
        dictArray = ["[OOV]"]
        for seq in texts:
            for word in seq:
                w = dictSet.get(word)
                if w == None:  # 如果当前词不在词典之中
                    dictSet[word] = len(dictArray)
                    dictArray.append(word)

        self.dict = (dictSet)
        self.dictArray = (dictArray)

        # 构造矩阵
        HMM.__init__(self, nClasses, len(dictArray))

    def getIndex(self, text :List[str]) -> List[int]:
        """把序列转换成索引"""
        indexes = []
        for word in text:
            index = self.dict.get(word)
            if(index == None):
                indexes.append(0)
            else:
                indexes.append(index)
        return indexes


    def getPI(self) -> np.ndarray:
        self.PI.fill(1)

        for labelSeq in self.labels:
            if(len(labelSeq) > 0):
                self.PI[labelSeq[0]] += 1

        self.PI /= self.PI.sum()

        return self.PI


    def getA(self) -> np.ndarray:
        self.A.fill(1)

        for labelSeq in self.labels:
            if(len(labelSeq) > 1):
                for i in range(1, len(labelSeq)):
                    labelPre = labelSeq[i - 1]
                    lableCur = labelSeq[i]
                    self.A[labelPre, lableCur] += 1

        #对每行进行归一化
        for i in range(self.N):
            maxValue = self.A[i,].sum()
            if(maxValue == 0):
                maxValue = 1
            self.A[i, ] /= maxValue

        return self.A


    def getB(self) -> np.ndarray:
        self.B.fill(1)
        if(len(self.texts) != len(self.labels)):
            raise RuntimeError("输入句子序列长度与输入标签序列长度不一致")

        for i in range(len(self.texts)):
            if (len(self.texts[i]) != len(self.labels[i])):
                raise RuntimeError("输入句子长度与标签长度不一致")
            for j in range(len(self.texts[i])) :
                state = self.labels[i][j]
                word = self.texts[i][j]
                wordIndex = self.dict[word]
                self.B[state, wordIndex] += 1

        for i in range(self.N):
            maxValue = self.B[i, ].sum()
            if(maxValue == 0):
                maxValue = 1
            self.B[i, ] /= maxValue

        return self.B



def viterbi(hmm:HMM, sequence: List[int]) -> List[int]:
    """ 把每个状态当成一个结点，结点的权重就是发射概率，结点的边就是转移概率，
        目标是找出具有最大概率的状态序列。
        我们可以使用动态规划的思想求解，即该问题具有最优子解构：
            如果长度为N的状态序列y1y2...yn是最优的，那么长度为N-1的状态序列y1y2..y(n-1)也必然是最优的！
        :param hmm  隐式马尔可夫对象
        :param sequence 输入序列
        :return 状态序列
     """

    # 处理特殊情况
    if(hmm == None or len(sequence) == 0):
        return []

    # 算法需要的数据结构
    scores = np.zeros(hmm.N, dtype=float)
    path = np.zeros((len(sequence), hmm.N), dtype=np.int)
    PI = hmm.getPI()
    A = hmm.getA()
    B = hmm.getB()

    # 算法开始
    # 使用对数，防止连乘越界
    for i in range(len(scores)):
        scores[i] = np.log(PI[i]) + np.log(B[i, sequence[0]])

    for i in range(1, len(sequence)):
        scoresTemp = np.ones(scores.shape) * (- np.inf)   #初始化为负无穷
        for s2 in range(hmm.N):         #后一个状态
            for s1 in range(hmm.N):     #前一个状态
                scoreCur = scores[s1] + np.log(A[s1, s2])  + np.log(B[s2, sequence[i]])
                if(scoreCur >= scoresTemp[s2]):
                    scoresTemp[s2] = scoreCur
                    path[i, s2] = s1
        scores = scoresTemp

    # 解最优路径
    bestEndStateIndex = scores.argmax()
    bestPath = [bestEndStateIndex]
    for i in range(1, len(sequence)):
        lastState = path[len(sequence) - i][bestEndStateIndex]
        bestPath.append(lastState)
        bestEndStateIndex = lastState

    bestPath.reverse()

    return bestPath